import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os

class AnomalyDataset(Dataset):
    def __init__(self, root, phase, transform=None, include_train_defect=False):
        """
        Args:
            root (str): 数据集根目录（包含train和test子文件夹）
            phase (str): 'train'或'test'
            transform (callable, optional): 图像变换
            include_train_defect (bool): 训练阶段是否包含异常样本（True=包含，False=不包含）
        """
        self.root = os.path.join(root, phase)  # 指向train或test子文件夹
        self.phase = phase
        self.transform = transform
        self.classes = {'normal': 0, 'defect': 1}  # 正常为0，缺陷为1
        self.include_train_defect = include_train_defect  # 新增参数：控制训练集是否加载异常样本
        self.samples = self._load_samples()

    def _load_samples(self):
        samples = []
        # 1. 加载正常样本（train和test都必须包含）
        normal_dir = os.path.join(self.root, 'normal_regions')
        if not os.path.exists(normal_dir):
            raise FileNotFoundError(f"正常样本目录不存在: {normal_dir}")
        for img_name in os.listdir(normal_dir):
            img_path = os.path.join(normal_dir, img_name)
            if os.path.isfile(img_path):  # 过滤非文件（如文件夹）
                samples.append((img_path, self.classes['normal']))

        # 2. 加载缺陷样本
        if self.phase == 'test':
            # 测试集必须包含缺陷样本
            defect_dir = os.path.join(self.root, 'defect_regions')
            if not os.path.exists(defect_dir):
                raise FileNotFoundError(f"测试集缺陷样本目录不存在: {defect_dir}")
            for img_name in os.listdir(defect_dir):
                img_path = os.path.join(defect_dir, img_name)
                if os.path.isfile(img_path):
                    samples.append((img_path, self.classes['defect']))
        elif self.phase == 'train' and self.include_train_defect:
            # 训练集：仅当include_train_defect=True时加载缺陷样本
            defect_dir = os.path.join(self.root, 'defect_regions')
            if not os.path.exists(defect_dir):
                raise FileNotFoundError(f"训练集缺陷样本目录不存在: {defect_dir}\n请确认是否需要在训练集加入异常样本")
            for img_name in os.listdir(defect_dir):
                img_path = os.path.join(defect_dir, img_name)
                if os.path.isfile(img_path):
                    samples.append((img_path, self.classes['defect']))

        return samples

    def __getitem__(self, idx):
        img_path, label = self.samples[idx]
        img = Image.open(img_path).convert('RGB')  # 转为RGB格式
        if self.transform:
            img = self.transform(img)
        return img, label  # 返回（图像，标签）

    def __len__(self):
        return len(self.samples)


def get_dataloader(root, batch_size, workers, isize, phase, include_train_defect=False):
    """
    获取DataLoader
    Args:
        include_train_defect (bool): 训练阶段是否包含异常样本（仅对phase='train'有效）
    """
    # 图像预处理（与原逻辑一致）
    transform = transforms.Compose([
        transforms.Resize(isize),
        transforms.CenterCrop(isize),
        transforms.ToTensor(),
        transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))  # 标准化到[-1, 1]
    ])

    # 初始化数据集（传入是否包含训练异常样本的参数）
    dataset = AnomalyDataset(
        root=root,
        phase=phase,
        transform=transform,
        include_train_defect=include_train_defect  # 传递参数
    )

    # 数据加载配置
    shuffle = phase == 'train'  # 训练集打乱，测试集不打乱
    dataloader = DataLoader(
        dataset,
        batch_size=batch_size,
        shuffle=shuffle,
        num_workers=workers,
        drop_last=phase == 'train'  # 训练集丢弃最后不完整批次
    )
    print("shuffle=", shuffle)
    return dataloader